# Copyright (c) SenseTime. All Rights Reserved.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from pysot.core.config import cfg
from pysot.models.loss import select_cross_entropy_loss, weight_l1_loss
from pysot.models.backbone import get_backbone
from pysot.models.head import get_rpn_head, get_mask_head, get_refine_head
from pysot.models.neck import get_neck

import matplotlib
from matplotlib import cm
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import scipy.io as scio
mat_seq = 0


class ModelBuilder(nn.Module):
    def __init__(self):
        super(ModelBuilder, self).__init__()
        # build backbone
        self.backbone = get_backbone(cfg.BACKBONE.TYPE,
                                     **cfg.BACKBONE.KWARGS)

        # build adjust layer
        if cfg.ADJUST.ADJUST:
            self.neck = get_neck(cfg.ADJUST.TYPE,
                                 **cfg.ADJUST.KWARGS)

        # build rpn head
        self.rpn_head = get_rpn_head(cfg.RPN.TYPE,
                                     **cfg.RPN.KWARGS)

        # build mask head
        if cfg.MASK.MASK:
            self.mask_head = get_mask_head(cfg.MASK.TYPE,
                                           **cfg.MASK.KWARGS)

            if cfg.REFINE.REFINE:
                self.refine_head = get_refine_head(cfg.REFINE.TYPE)

    def template(self, z):
        zf = self.backbone(z)
        if cfg.MASK.MASK:
            zf = zf[-1]
        if cfg.ADJUST.ADJUST:
            zf = self.neck(zf)
        self.zf = zf


    ###############################################################
    # get top distance:
    def get_top(self, arr, top_k):
        brr = arr.ravel()
        top_k_index = brr.argsort()[::-1][0:top_k]
        ori_index = []
        for ind in top_k_index:
            row = ind // 25
            col = ind % 25
            ori_index.append((row, col))

        return ori_index

    def remove_low_point(self, top_k, score):
        top_1 = top_k[0]
        top_1_point = score[top_1[0]][top_1[1]]
        new_top_4 = []
        new_top_4.append(top_1)
        for iqq in range(1, len(top_k)):
            if (score[top_k[iqq][0]][top_k[iqq][1]] > top_1_point * 0.75):
                new_top_4.append(top_k[iqq])
        return new_top_4

    def is_neighbour(self, ori_pos, insert_pos):
        if (insert_pos[0] <= ori_pos[0] + 2 and insert_pos[0] >= ori_pos[0] - 2 and insert_pos[1] <= ori_pos[1] + 2 and
                insert_pos[1] >= ori_pos[1] - 2):
            return True
        return False

    def check_neighbour(self, ori_arr, insert_pos):
        for arr_insex in range(len(ori_arr)):
            ff_flag = self.is_neighbour(ori_arr[arr_insex], insert_pos)
            if (ff_flag == True):
                return True
        return False

    def merge_neighbour(self, arr_pos):
        new_list = []
        new_list.append(arr_pos[0])
        for index in range(1, len(arr_pos)):
            check_flag = self.check_neighbour(new_list, arr_pos[index])
            if (check_flag == False):
                new_list.append(arr_pos[index])
        return new_list

    def compute_dis(self, ulti_top):
        if (len(ulti_top) < 2):
            return 10
        else:
            distance_arr = []
            for i in range(1, len(ulti_top)):
                n_pos = ulti_top[0]
                temp_pos = ulti_top[i]
                dis = np.sqrt((temp_pos[0] - n_pos[0]) * (temp_pos[0] - n_pos[0]) + (temp_pos[1] - n_pos[1]) * (
                            temp_pos[1] - n_pos[1]))
                distance_arr.append(dis)
            return distance_arr
    ###############################################################



    def track(self, x):
        xf = self.backbone(x)
        if cfg.MASK.MASK:
            self.xf = xf[:-1]
            xf = xf[-1]
        if cfg.ADJUST.ADJUST:
            xf = self.neck(xf)


        ###############################################################
        # # heat map:
        # x_x = F.conv2d(xf[0], self.zf[0])
        # tmp = x_x[0, 0, :, :]
        # vector_output = tmp.cpu().detach().numpy()
        # fig = plt.figure()
        # fig.set_size_inches(2, 2)
        # fig.canvas.manager.window.move(500, 300)
        # plt.imshow(vector_output)
        # plt.axis('off')
        #
        # plt.gca().xaxis.set_major_locator(plt.NullLocator())
        # plt.gca().yaxis.set_major_locator(plt.NullLocator())
        # plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
        # plt.margins(0, 0)
        # plt.pause(1.0)
        # plt.close(fig)
        ###############################################################


        ###############################################################
        # # save .mat:
        # global mat_seq
        # mat_seq = mat_seq + 1
        # x_x = F.conv2d(xf[0],self.zf[0])
        # tmp = x_x[0,0,:,:]
        # vector_output = tmp.cpu().detach().numpy()
        # fig = plt.figure()
        # fig.set_size_inches(2,2)
        # fig.canvas.manager.window.move(500,300)
        # plt.imshow(vector_output)
        # plt.axis('off')
        #
        # model_path = 'hotmap/girl/'
        # if not os.path.isdir(model_path):
        #     os.makedirs(model_path)
        # sv_mat_path = os.path.join(model_path, '{}.txt'.format(mat_seq))
        # # sv_mat_path = 'hotmap/girl/{:08}.png'.format(mat_seq)
        #
        # # scio.savemat(sv_mat_path, {'vector:': vector_output})
        # plt.gca().xaxis.set_major_locator(plt.NullLocator())
        # plt.gca().yaxis.set_major_locator(plt.NullLocator())
        # plt.subplots_adjust(top=1,bottom=0,left=0,right=1,hspace=0,wspace=0)
        # plt.margins(0,0)
        #
        # model_path = 'hotmap/girl/'
        # if not os.path.isdir(model_path):
        #     os.makedirs(model_path)
        # plt.savefig(os.path.join(model_path, '{:08}.png'.format(mat_seq)))
        # # plt.savefig('hotmap/girl/{:08}.png'.format(mat_seq))
        #
        # plt.pause(1.0)
        # plt.close(fig)
        ###############################################################



        cls, loc = self.rpn_head(self.zf, xf)
        if cfg.MASK.MASK:
            mask, self.mask_corr_feature = self.mask_head(self.zf, xf)


        ###############################################################
        # get top distance:
        # print(len(xf))
        # print(len(self.zf))
        # backend = matplotlib.get_backend()
        # print(backend)

        x_tmp = xf[1]
        z_tmp = self.zf[1]
        feature_tmp = F.conv2d(x_tmp, z_tmp)
        vector_tmp = feature_tmp[0, 0, :, :].cpu().detach().numpy()
        top_k = self.get_top(vector_tmp, 50)
        top_k_1 = self.merge_neighbour(top_k)
        top_k_2 = self.remove_low_point(top_k_1, vector_tmp)
        top_k_3 = self.compute_dis(top_k_2)
        # names = locals()
        # for i in range(3):
        #     x_tmp = xf[i]
        #     z_tmp = self.zf[i]
        #     feature_tmp = F.conv2d(x_tmp,z_tmp)
        #     vector_tmp = feature_tmp[0,0,:,:].cpu().detach().numpy()
        #     # print(vector_tmp.shape)
        x_arr = np.arange(0, 25, 1)
        y_arr = np.arange(0, 25, 1)
        x_arr, y_arr = np.meshgrid(x_arr, y_arr)
        z_arr = np.zeros((25, 25))
        for i1 in range(25):
            for j1 in range(25):
                z_arr[i1][j1] = vector_tmp[i1][j1]
        fig_tmp = plt.figure()
        #     names['fig_'+str(i+1)] = plt.figure()
        pos = '+500+300'
        # fig_tmp.canvas.manager.window.wm_geometry(pos)
        fig_tmp.canvas.manager.window.move(500, 300)
        ax = Axes3D(fig_tmp)
        ax.plot_surface(x_arr, y_arr, z_arr, cmap=cm.coolwarm)

        # # show Axes3D
        # plt.pause(0.5)
        # plt.close(fig_tmp)

        # plt.close(names['fig_'+str(2)])
        # plt.close(names['fig_'+str(3)])
        ###############################################################


        return {
                'cls': cls,
                'loc': loc,
                'mask': mask if cfg.MASK.MASK else None,
                'distance': top_k_3
               }

    def mask_refine(self, pos):
        return self.refine_head(self.xf, self.mask_corr_feature, pos)

    def log_softmax(self, cls):
        b, a2, h, w = cls.size()
        cls = cls.view(b, 2, a2//2, h, w)
        cls = cls.permute(0, 2, 3, 4, 1).contiguous()
        cls = F.log_softmax(cls, dim=4)
        return cls

    def forward(self, data):
        """ only used in training
        """
        template = data['template'].cuda()
        search = data['search'].cuda()
        label_cls = data['label_cls'].cuda()
        label_loc = data['label_loc'].cuda()
        label_loc_weight = data['label_loc_weight'].cuda()

        # get feature
        zf = self.backbone(template)
        xf = self.backbone(search)
        if cfg.MASK.MASK:
            zf = zf[-1]
            self.xf_refine = xf[:-1]
            xf = xf[-1]
        if cfg.ADJUST.ADJUST:
            zf = self.neck(zf)
            xf = self.neck(xf)
        cls, loc = self.rpn_head(zf, xf)

        # get loss
        cls = self.log_softmax(cls)
        cls_loss = select_cross_entropy_loss(cls, label_cls)
        loc_loss = weight_l1_loss(loc, label_loc, label_loc_weight)

        outputs = {}
        outputs['total_loss'] = cfg.TRAIN.CLS_WEIGHT * cls_loss + \
            cfg.TRAIN.LOC_WEIGHT * loc_loss
        outputs['cls_loss'] = cls_loss
        outputs['loc_loss'] = loc_loss

        if cfg.MASK.MASK:
            # TODO
            mask, self.mask_corr_feature = self.mask_head(zf, xf)
            mask_loss = None
            outputs['total_loss'] += cfg.TRAIN.MASK_WEIGHT * mask_loss
            outputs['mask_loss'] = mask_loss
        return outputs
